import torch
from torch.utils.data import Dataset
import numpy as np

def from_npz(npz_path):
    npz = np.load(npz_path)
    arr = npz['arr_0']
    return arr

class CBDataset(Dataset):
    def __init__(self, dts=[], data=None):
        self.source = dts
        self.data = data

    def load(self):
        for each in self.source:
            if each.endswith('npz'):
                arr = from_npz(each)
                # print(f'arr: {arr.shape}')
                if self.data is None:
                    self.data = arr
                else:
                    self.data = np.vstack((self.data, arr))
            else:
                print(each)
                raise NotImplementedError
        return self
    
    def to_tensor(self, device):
        self.data = torch.tensor(self.data, device=device)
        return self
            
    def __len__(self):
        return len(self.data) if self.data is not None else 0
    
    def __getitem__(self, index):
        inputs = self.data[index, :5, :] # (,5,5)
        target = self.data[index, 5, :3] # (,,3)
        if not isinstance(inputs, torch.Tensor):
            inputs = torch.tensor(inputs)
        if not isinstance(target, torch.Tensor):
            target = torch.tensor(target)
        return inputs, target

def create_balanced_sampler(dataset, number):
    # 使用PyTorch的方法计算最大值
    max_stims = torch.max(dataset.data[:,:5,3:], dim=1)[0].max(dim=1)[0]
    # 使用PyTorch的方法查找索引
    zero_indices = torch.where(max_stims == 0)[0]
    nonzero_indices = torch.where(max_stims > 0)[0]
    class BalancedSampler(torch.utils.data.Sampler):
        def __init__(self, zero_indices, nonzero_indices, samples_per_epoch=1000):
            self.zero_indices = zero_indices
            self.nonzero_indices = nonzero_indices
            self.samples_per_epoch = samples_per_epoch
        def __iter__(self):
            if len(self.zero_indices) > self.samples_per_epoch:
                # 使用PyTorch的随机采样方法代替NumPy的random.choice
                indices = torch.randperm(len(self.zero_indices))[:self.samples_per_epoch]
                sampled_zero_indices = self.zero_indices[indices]
            else:
                sampled_zero_indices = self.zero_indices.clone()
                # sampled_zero_indices = np.random.choice(self.zero_indices, self.samples_per_epoch, replace=True)
            # 使用PyTorch的cat方法代替NumPy的concatenate
            indices = torch.cat([sampled_zero_indices, self.nonzero_indices])
            # 使用PyTorch的方法打乱索引
            indices = indices[torch.randperm(len(indices))]
            return iter(indices.tolist())
        def __len__(self):
            return len(self.nonzero_indices) + self.samples_per_epoch
    return BalancedSampler(zero_indices, nonzero_indices, number)

if __name__ == '__main__':
    npz = from_npz(r"C:\Users\songy\Desktop\data-0518\fyp_20250331120047\ds.npz")
    